from tools.data_gen import data_flow
from model.train_model import train_model
import tensorflow as tf
from model.predict_model import init_artificial_neural_network, prediction_result_from_img
from model.test_model import test_model
import matplotlib.pyplot as plt


# 全局配置文件
tf.compat.v1.app.flags.DEFINE_string('test_data_local', './test_image', '测试图片文件夹')
tf.compat.v1.app.flags.DEFINE_string('data_local', './jupyterLab/data/base_data', '训练图片文件夹')
tf.compat.v1.app.flags.DEFINE_integer('num_classes', 40, '垃圾分类数目')
tf.compat.v1.app.flags.DEFINE_integer('input_size', 224, '模型输入图片大小')
tf.compat.v1.app.flags.DEFINE_integer('batch_size', 16, '图片批处理大小')
tf.compat.v1.app.flags.DEFINE_float('learning_rate', 1e-4, '学习率')
tf.compat.v1.app.flags.DEFINE_integer('max_epochs', 4, '步长')
tf.compat.v1.app.flags.DEFINE_string('train_local', './output_model/', '训练输出文件夹')
tf.compat.v1.app.flags.DEFINE_integer('keep_weights_file_num', 20, '如果设置为-1，则文件保持的最大权重数表示无穷大')
FLAGS = tf.compat.v1.app.flags.FLAGS


def train(train_sequence, validation_sequence):
    history_tl_50, history_tl = train_model(FLAGS, train_sequence, validation_sequence)
    draw(history_tl_50)
    draw(history_tl)


def predict():
    model = init_artificial_neural_network(FLAGS)
    while True:
        try:
            img_url = input("请输入图片地址:")
            print('您输入的图片地址为：' + img_url)
            res = prediction_result_from_img(model, img_url, FLAGS)
        except Exception as e:
            print('发生了异常：', e)


def test_res():
    model = init_artificial_neural_network(FLAGS)
    test_acc = test_model(test_img_paths, test_labels, model, FLAGS)
    print('test_acc is: ', test_acc)


def draw(history):
    history_dict = history.history
    loss_values = history_dict['loss']
    val_loss_values = history_dict['val_loss']
    epochs = range(1, len(loss_values) + 1)

    # Training and validation loss
    plt.plot(epochs, loss_values, 'bo', label='Training loss')
    plt.plot(epochs, val_loss_values, 'b', label='Validation loss')
    plt.title('Training and validation loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    plt.show()

    # Training and validation acc
    acc = history_dict['acc']
    val_acc = history_dict['val_acc']

    plt.plot(epochs, acc, 'bo', label='Training acc')
    plt.plot(epochs, val_acc, 'b', label='Validation acc')
    plt.title('Training and validation acc')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.legend()

    plt.show()


if __name__ == '__main__':
    train_sequence, validation_sequence, test_img_paths, test_labels = data_flow(FLAGS.data_local, FLAGS.batch_size,
                                                                                 FLAGS.num_classes, FLAGS.input_size)
    test_res()